package com.kasax.plugin.app.transform

import com.android.build.api.transform.*
import com.android.build.gradle.internal.pipeline.TransformManager
import com.kasax.plugin.app.LogUtils
import com.kasax.plugin.app.visitor.moduleinit.MooduleInitClassVisitor
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter

import java.nio.file.Files
import java.nio.file.Path
import java.util.jar.JarEntry
import java.util.jar.JarFile

class SubModuleInitTransform extends Transform {
    final String TAG = "SubModuleInitTransform"

    /**
     * 当前Transform的名称
     * @return
     */
    @Override
    String getName() {
        return "SubModuleInitTransform"
    }
    /**
     * 返回告知编译器，当前Transform需要消费的输入类型
     * @return
     */
    @Override
    Set<QualifiedContent.ContentType> getInputTypes() {
//        LogUtils.d(TAG," getInputTypes")
        return TransformManager.CONTENT_CLASS
    }
    /**
     * 告知编译器，当前Transorm需要收集的范围
     * @return
     */
    @Override
    Set<? super QualifiedContent.Scope> getScopes() {
//        LogUtils.d(TAG," getScopes")
        return [
                QualifiedContent.Scope.PROJECT,
                QualifiedContent.Scope.SUB_PROJECTS,
                QualifiedContent.Scope.EXTERNAL_LIBRARIES
        ]
    }
    /**
     * 是否支持增量
     * @return
     */
    @Override
    boolean isIncremental() {
//        LogUtils.d(TAG," isIncremental")
        return false
    }
    /**
     * 所有的Class收集好以后，会被打包进此方法
     * @param transformInvocation the invocation object containing the transform inputs.
     * @throws TransformException* @throws InterruptedException* @throws IOException
     */
    @Override
    void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        LogUtils.d(TAG, "transform")
        LogUtils.d(TAG, "Inputs count: ${transformInvocation.getInputs().size()}")
        for (TransformInput input : transformInvocation.getInputs()) {
            LogUtils.d(TAG, "DirectoryInputs count: ${input.getDirectoryInputs().size()}")
            LogUtils.d(TAG, "JarInputs count: ${input.getJarInputs().size()}")
            for (DirectoryInput directoryInput : input.getDirectoryInputs()) {
                LogUtils.d(TAG, "Processing directory: ${directoryInput.getFile().absolutePath}")
                LogUtils.d(TAG, "DirectoryInput name: ${directoryInput.name}")
                processDirectory(directoryInput.getFile(), transformInvocation.getOutputProvider())
            }

            for (JarInput jarInput : input.getJarInputs()) {
                LogUtils.d(TAG, "Processing jar: ${jarInput.getFile().absolutePath}")
                LogUtils.d(TAG, "JarInput name: ${jarInput.name}")
                processJar(jarInput.getFile(), transformInvocation.getOutputProvider())
            }
        }
    }

    void processDirectory(File sourceFile, TransformOutputProvider outputProvider) {
        File outputDir = outputProvider.getContentLocation("SubModuleInitTransformOutput", getOutputTypes(), getScopes(), Format
                .DIRECTORY)
        LogUtils.d(TAG, " processDirectory sourceFile: ${sourceFile.absolutePath} , outputDir: ${outputDir.absolutePath}")
        if (sourceFile.isDirectory() && sourceFile.exists()) {
            processDIrecotryRecursively(sourceFile, outputDir)
        }
    }

    void processDIrecotryRecursively(File sourceFile, File outputFile) {
        // 使用 listFiles 替代 Files.walk
        File[] files = sourceFile.listFiles()
        if (files == null || files.size() == 0) {
            return
        }

        for (File file : files) {
            if (file.isDirectory()) {
                // 递归处理文件
                processDIrecotryRecursively(file, outputFile)
            } else {
                // 处理class文件
                processClassFile(file, outputFile)
            }
        }
    }

    void processJar(File sourceFile, TransformOutputProvider outputProvider) {
        File outputDir = outputProvider.getContentLocation("SubModuleInitTransformOutput", getOutputTypes(), getScopes(), Format.DIRECTORY)
        LogUtils.d(TAG, " processJar sourceFile: ${sourceFile.absolutePath} , outputDir: ${outputDir.absolutePath}")

        if (!sourceFile.exists()) {
            LogUtils.d(TAG, "Skipping non-existent jar: ${sourceFile.absolutePath}")
            return
        }

        // Ensure output directory exists
        if (!outputDir.exists()) {
            LogUtils.d(TAG, "Creating output directory: ${outputDir.absolutePath}")
            outputDir.mkdirs()
        }

        try {
            // Extract and process classes from the JAR file
            JarFile jarFile = new JarFile(sourceFile)
            Enumeration<JarEntry> entries = jarFile.entries()

            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement()
                if (entry.isDirectory() || !entry.getName().endsWith('.class')) {
                    continue
                }

                // Read the class file from the JAR
                InputStream inputStream = jarFile.getInputStream(entry)
                byte[] classBytes = inputStream.readAllBytes()
                inputStream.close()

                // Modify the class file
                byte[] modifiedBytes = modifyClass(classBytes)

                // Write the modified class to the output directory
                String entryName = entry.getName().replace('/', File.separator)
                File outputFile = new File(outputDir, entryName)
                outputFile.getParentFile().mkdirs()
                Files.write(outputFile.toPath(), modifiedBytes)
                LogUtils.d(TAG, "Processed and saved class: ${entryName}")
            }

            jarFile.close()
        } catch (Exception e) {
            LogUtils.d(TAG, "Failed to process jar ${sourceFile.absolutePath}: ${e.message}")
            throw e
        }
    }

    void processClassFile(File sourceFile, File outputDir) {
        if (!sourceFile.exists()) {
            LogUtils.d(TAG, "Skipping non-existent file: ${sourceFile.absolutePath}")
            return
        }
        LogUtils.d(TAG, " processClassFile sourceFile: ${sourceFile.absolutePath} , outputDir: ${outputDir.absolutePath}")

        // Ensure output directory exists
        if (!outputDir.exists()) {
            LogUtils.d(TAG, "Creating output directory: ${outputDir.absolutePath}")
            if (!outputDir.mkdirs()) {
                LogUtils.d(TAG, "Failed to create output directory: ${outputDir.absolutePath}")
                return
            }
        }

        try {
            byte[] classBytes = Files.readAllBytes(sourceFile.toPath())
            LogUtils.d(TAG, "Processing file: ${sourceFile.absolutePath}, size: ${classBytes.length} bytes")
            byte[] modifiedBytes = modifyClass(classBytes)
            Path outputPath = outputDir.toPath().resolve(sourceFile.getName())
            Files.write(outputPath, modifiedBytes)
            LogUtils.d(TAG, "Successfully processed and saved file: ${outputPath.toAbsolutePath()}")
        } catch (Exception e) {
            LogUtils.d(TAG, "Failed to process file ${sourceFile.absolutePath}: ${e.message}")
        }
    }

    byte[] modifyClass(byte[] originalBytes) {
        ClassReader clsReader = new ClassReader(originalBytes)
        ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS)
        MooduleInitClassVisitor classVisitor = new MooduleInitClassVisitor(classWriter)
        clsReader.accept(classVisitor, 0)
        return classWriter.toByteArray()
    }
}